# -*-coding:utf-8-*-
from Basic import *
import pytesseract as pytesseract
from PIL import Image, ImageDraw
from selenium.webdriver import ActionChains

class ImgHandle():
    ''' 验证码灰处理，降噪  '''
    def blackimg(self, img_path):
        ''' 图像灰处理  '''
        img = Image.open(img_path)
        w, h = img.size
        for x in range(w):
            for y in range(h):
                r, g, b = img.getpixel((x, y))
                if 190 <= r <= 255 and 170 <= g <= 255 and 0 <= b <= 140:
                    img.putpixel((x, y), (0, 0, 0))
                if 0 <= r <= 90 and 210 <= g <= 255 and 0 <= b <= 90:
                    img.putpixel((x, y), (0, 0, 0))
        img = img.convert('L').point([0] * 150 + [1] * (256 - 150), '1')
        return img

    def twoValue(self, image, G):
        # 二值数组
        self.t2val = {}
        for y in range(0, image.size[1]):
            for x in range(0, image.size[0]):
                g = image.getpixel((x, y))
                if g > G:
                    self.t2val[(x, y)] = 1
                else:
                    self.t2val[(x, y)] = 0

    # 根据一个点A的RGB值，与周围的8个点的RBG值比较，设定一个值N（0 <N <8），当A的RGB值与周围8个点的RGB相等数小于N时，此点为噪点
    # G: Integer 图像二值化阀值
    # N: Integer 降噪率 0 <N <8
    # Z: Integer 降噪次数
    # 输出
    #  0：降噪成功
    #  1：降噪失败
    def clearNoise(self, image, N, Z):
        ''' 验证码降噪  '''
        for i in range(0, Z):
            self.t2val[(0, 0)] = 1
            self.t2val[(image.size[0] - 1, image.size[1] - 1)] = 1
            for x in range(1, image.size[0] - 1):
                for y in range(1, image.size[1] - 1):
                    nearDots = 0
                    L = self.t2val[(x, y)]
                    if L == self.t2val[(x - 1, y - 1)]:
                        nearDots += 1
                    if L == self.t2val[(x - 1, y)]:
                        nearDots += 1
                    if L == self.t2val[(x - 1, y + 1)]:
                        nearDots += 1
                    if L == self.t2val[(x, y - 1)]:
                        nearDots += 1
                    if L == self.t2val[(x, y + 1)]:
                        nearDots += 1
                    if L == self.t2val[(x + 1, y - 1)]:
                        nearDots += 1
                    if L == self.t2val[(x + 1, y)]:
                        nearDots += 1
                    if L == self.t2val[(x + 1, y + 1)]:
                        nearDots += 1
                    if nearDots < N:
                        self.t2val[(x, y)] = 1

    def saveImage(self, img_path, size):
        ''' 保存图片  '''
        image = Image.new("1", size)
        draw = ImageDraw.Draw(image)
        for x in range(0, size[0]):
            for y in range(0, size[1]):
                draw.point((x, y), self.t2val[(x, y)])
        image.save(img_path)

    def img_pcocessing(self, img_path):
        ''' 验证码灰处理，降噪  '''
        image = self.blackimg(img_path)
        #图片改名并保存
        img_path = img_path.replace('jpg', 'png')
        image.save(img_path)

        #图像降噪 并改为jpg ,只支持png
        image = Image.open(img_path).convert("L")
        self.twoValue(image, 100)
        self.clearNoise(image, 3, 2)

        # code_path = os.path.join(log_path,"识别图片.jpeg")
        self.saveImage(img_path, image.size)
        return img_path


    def get_imgnumber(self,img_path):
        ''' 验证码识别主程序  '''
        img_path = self.img_pcocessing(img_path)#灰处理降噪
        ''' 验证码识别模块  '''
        im = Image.open(img_path)
        num = pytesseract.image_to_string(im)
        # strs = res.split("\n")
        return num.replace(" ", "")


# class GetCodeNum():

    def IsValidImage(self,img_path):
        """
        判断文件是否为有效（完整）的图片
        :param img_path:图片路径
        :return:True：有效 False：无效
        """
        bValid = True
        try:
            Image.open(img_path).verify()
        except:
            bValid = False
        return bValid

    def transjpg(self,img_path):
        """转换图片格式 """
        if self.IsValidImage(img_path):
            str = img_path.rsplit(".", 1)
            output_imgjpg_path = str[0] + ".jpg"
            # print(output_img_path)
            im = Image.open(img_path)
            im = im.convert('RGB') #jpg不支持RGB，需去除
            im.save(output_imgjpg_path)
            return output_imgjpg_path


    def ele_screenshot(self, driver, ele, img_path):
        ''' 获取验证码结果
           driver  ：seleium实例
           ele     ：定位到的元素
           img_path：图片保存路径
        '''
        if  'jpg' in img_path:
            img_path = img_path.replace("jpg", "png")
            # print('int png_imgpath：%s' % imgpath)

        driver.save_screenshot(img_path)  # 截取全屏
        # print(ele.location)  #打印元素坐标
        # print(ele.size)  #获取大小如： {'height': 50, 'width': 120}
        left = ele.location['x']
        top = ele.location['y']
        right = ele.location['x'] + ele.size['width']
        bottom = ele.location['y'] + ele.size['height']
        im = Image.open(img_path)
        im = im.crop((left, top, right, bottom))
        # imgpath = imgpath.replace("png", "jpg")
        im.save(img_path)

        #将图片格式改为jpg
        imgjpg_path = self.transjpg(img_path)
        # print('out jpg_imgpath：%s' % imgpath)
        return imgjpg_path

    def get_ele_code(self,driver,ele,img_path,imgtxtlen=4):
        '''获取验证码结果
           driver  ：seleium实例
           ele     ：定位到的元素
           img_path：图片保存路径
        '''
        imgtext = 0
        imgflag = 1
        num = 0
        while imgflag :
            num = num+1
            if os.path.exists(img_path):  # 如果文件存在
                # print('--   --> %s 检测到图片存在，删除该文件：%s ' % (num,code_path))
                os.remove(img_path)

            # print('--   --> ①截全屏，②获取元素大小及位置，③并截取元素图片：%s ' % code_path)
            img_path = self.ele_screenshot(driver,ele,img_path)#截取元素图片
            # code_path = imgP.transjpg(code_path) #修改文件格式
            imgtext = self.get_imgnumber(img_path) #进行图像识别

            print('--   --> 验证码识别结果：%s' % imgtext)

            if len(imgtext) == imgtxtlen:
                # print('--   --> %s 识别出的验证码为四个字符：%s' % (num,imgtext))
                imgflag = 0
            else:
                # print('--   --> 点击刷新验证码')
                action = ActionChains(driver).move_to_element(ele)  # 移动到该元素
                action.click(ele).perform()
        return imgtext


if __name__ == '__main__':
    from selenium import webdriver
    img = ImgHandle()
    driver = webdriver.Chrome()
    driver.maximize_window()
    driver.get('https://192.168.63.192:8443/dolphin')
    ele = driver.find_element_by_xpath('//*[@id="login_yzm_img"]')
    # img.ele_screenshot(driver,ele,code_path)
    img.get_ele_code(driver,ele,os.path.join(os.getcwd(),"getcode.jpg"))
    driver.close()


